package net.cyue.ort.llm.sampling;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 采样辅助工具类
 * 抽取常见的logit处理逻辑，降低具体策略之间的耦合
 */
public final class SamplingUtils {
    private SamplingUtils() {
    }

    public static double[] toDoubleArray(float[] logits) {
        if (logits == null) {
            return new double[0];
        }
        double[] result = new double[logits.length];
        for (int i = 0; i < logits.length; i++) {
            result[i] = logits[i];
        }
        return result;
    }

    public static void applyTemperature(double[] logits, float temperature) {
        if (logits.length == 0 || temperature <= 0f || temperature == 1f) {
            return;
        }
        double temp = temperature;
        for (int i = 0; i < logits.length; i++) {
            logits[i] /= temp;
        }
    }

    public static void applyRepetitionPenalty(double[] logits, List<Long> context, float penalty) {
        if (logits.length == 0 || context == null || context.isEmpty() || penalty <= 1.0f) {
            return;
        }
        for (Long tokenId : context) {
            if (tokenId == null) {
                continue;
            }
            int idx = tokenId.intValue();
            if (idx < 0 || idx >= logits.length) {
                continue;
            }
            double value = logits[idx];
            logits[idx] = value >= 0 ? value / penalty : value * penalty;
        }
    }

    public static List<Candidate> buildCandidates(double[] logits, int topK, float topP) {
        int vocabSize = logits.length;
        if (vocabSize == 0) {
            return List.of();
        }

        Integer[] indices = new Integer[vocabSize];
        Arrays.setAll(indices, i -> i);
        Arrays.sort(indices, (a, b) -> Double.compare(logits[b], logits[a]));

        int limit = topK > 0 ? Math.min(topK, vocabSize) : vocabSize;
        double maxLogit = Double.NEGATIVE_INFINITY;
        for (int i = 0; i < limit; i++) {
            maxLogit = Math.max(maxLogit, logits[indices[i]]);
        }

        double[] expValues = new double[limit];
        double expSum = 0d;
        for (int i = 0; i < limit; i++) {
            double val = Math.exp(logits[indices[i]] - maxLogit);
            expValues[i] = val;
            expSum += val;
        }

        double threshold = (topP > 0f && topP < 1f) ? topP : 1d;
        List<Candidate> candidates = new ArrayList<>();
        double cumulative = 0d;
        for (int i = 0; i < limit; i++) {
            double probability = expValues[i] / expSum;
            cumulative += probability;
            candidates.add(new Candidate(indices[i], probability, logits[indices[i]]));
            if (cumulative >= threshold && !candidates.isEmpty()) {
                break;
            }
        }

        if (candidates.isEmpty()) {
            candidates.add(new Candidate(indices[0], expValues[0] / expSum, logits[indices[0]]));
        }

        return candidates;
    }

    public static final class Candidate {
        private final int tokenId;
        private final double probability;
        private final double logit;

        public Candidate(int tokenId, double probability, double logit) {
            this.tokenId = tokenId;
            this.probability = probability;
            this.logit = logit;
        }

        public int tokenId() {
            return tokenId;
        }

        public double probability() {
            return probability;
        }

        public double logit() {
            return logit;
        }
    }
}

